library(here)
library(data.table)
library(PanelMatch)

fullData <- fread(here('data', 'full_rol.csv'))

# Estimate alternative cutoffs ------------------------------------------------------

# This function takes about 10 minutes on a 2020 MacBook Pro
# The results are stored in 'data/alt_cutoff.csv'

alt_cutoff_results <- data.frame()

set.seed(10)
for (c in seq(.50, .9, .01)) {
  
  fullData$low_rol <- 1
  
  for (r in which(fullData$v2x_rule > c & fullData$ml_trans == 1)) {
    country <- fullData[r, iso3c]
    fullData[iso3c == country & ml == 1, low_rol := 0]
  }
  
  analysisData <- as.data.frame(fullData[ , .(year, cid_pm, ml, v2x_rule,
                                              bits_ln, market_ln, gdppc_ln,
                                              trade_ln, fdi_stock_ln, growth,
                                              nyc, low_rol)])
  
  pm_results <- PanelMatch(lag = 5, lead = 0:5,
                           time.id = "year", 
                           unit.id = "cid_pm", 
                           treatment = "ml", 
                           outcome.var = "v2x_rule",
                           refinement.method = "CBPS.weight", 
                           covs.formula = ~ lag(v2x_rule, 1) +
                             lag(bits_ln, 1) +
                             lag(market_ln, 1) +
                             lag(gdppc_ln, 1) +
                             lag(trade_ln, 1) +
                             lag(fdi_stock_ln, 1) +
                             lag(growth, 1) +
                             lag(nyc, 1),
                           qoi = "att", 
                           use.diagonal.variance.matrix = TRUE,
                           forbid.treatment.reversal = TRUE,
                           data = analysisData[analysisData$low_rol == 1, ])
  
  pe <- PanelEstimate(pm_results, number.iterations = 5000, 
                      data = analysisData[analysisData$low_rol == 1,])
  results <- summary(pe)$summary
  
  # Adds 90% CIs
  ci90 <- t(apply(pe$bootstrapped.estimates, 2, quantile, probs = c(.05, .95)))
  
  alt_cutoff_results <- rbind(alt_cutoff_results,
                              data.frame(cutoff = c,
                                         t = 0:5,
                                         number_cases = length(pe$matched.sets),
                                         estimate = results[ , 1],
                                         se = results[ , 2],
                                         ci_lower_95 = results[ , 3],
                                         ci_upper_95 = results[ , 4],
                                         ci_lower_90 = ci90[ , 1],
                                         ci_upper_90 = ci90[ , 2],
                                         asl = apply(pe$bootstrapped.estimates, 
                                                     2, 
                                                     function (x) sum(x > 0) / length(x))))
  
  print(c)
}
rm(c, pe, pm_results, analysisData, results, country, r, ci90)

setDT(alt_cutoff_results)

fwrite(alt_cutoff_results, here('data', 'alt_cutoff.csv'))

alt_cutoff_results$t <- factor(alt_cutoff_results$t, levels = 0:5)

alt_cutoff_results$t_name <- paste0('t=',alt_cutoff_results$t)

ggplot(alt_cutoff_results[as.integer(t) %in% c(1,3,6)],
       aes(x = cutoff, y = estimate)) +
  geom_hline(yintercept = 0, size = .3,
             color = 'red', linetype = 'dashed') +
  geom_line() +
  geom_ribbon(aes(ymax = ci_upper_95, ymin = ci_lower_95),
              alpha = .2) +
  geom_ribbon(aes(ymax = ci_upper_90, ymin = ci_lower_90),
              alpha = .35) +
  labs(y = 'ATT',
       x = 'Low Rule of Law Cutoff') +
  theme_bw(base_size = 11) +
  theme(panel.grid = element_blank()) +
  facet_wrap(~t_name)

ggsave(here('output', 'figures', 'figure-6.pdf'),
       height = 3, width = 8, units = 'in')

ggplot(alt_cutoff_results,
       aes(x = cutoff, y = estimate)) +
  geom_hline(yintercept = 0, size = .3,
             color = 'red', linetype = 'dashed') +
  geom_line() +
  geom_ribbon(aes(ymax = ci_upper_95, ymin = ci_lower_95),
              alpha = .2) +
  geom_ribbon(aes(ymax = ci_upper_90, ymin = ci_lower_90),
              alpha = .35) +
  labs(y = 'ATT',
       x = 'Low Rule of Law Cutoff') +
  theme_bw(base_size = 11) +
  theme(panel.grid = element_blank()) +
  facet_wrap(~t_name)

ggsave(here('output', 'figures', 'figure-A4.pdf'),
       height = 6, width = 8, units = 'in')
